/*
 * Copyright 2006-2007 S W McIlroy
 * Licensed  under the  Apache License,  Version 2.0  (the "License");
 * you may not use  this file  except in  compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed  under the  License is distributed on an "AS IS" BASIS,
 * WITHOUT  WARRANTIES OR CONDITIONS  OF ANY KIND, either  express  or
 * implied.
 *
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.mockfu;

import static org.mockfu.Exceptions.*;
import static org.objectweb.asm.Opcodes.*;

import java.lang.reflect.Constructor;
import java.lang.reflect.Field;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.lang.reflect.Modifier;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

import org.objectweb.asm.ClassAdapter;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.Label;
import org.objectweb.asm.Type;
import org.objectweb.asm.commons.GeneratorAdapter;
import org.objenesis.ObjenesisStd;

/**
 * A <tt>ProxyBuilder</tt> generates dynamic proxies that implement the non-final interface of any
 * class, a subclass whose method implementation is handled by a <tt>{@link MethodInterceptor}</tt>.
 *
 * @param <T>
 */
public class ProxyBuilder<T>
{
    static final String SUPER_METHOD_SUFFIX = "$$Super$$";
    static final String INTERCEPTOR_FIELD = "$$interceptor$$";

    /* tracks the classes that are currently being generated, so the same class isn't generated twice */
    private static final ConcurrentMap<String, GenerationLock> generationLocks
            = new ConcurrentHashMap<String, GenerationLock>();

    /**
     * A lock that waits for a particular class to be generated.
     */
    private static class GenerationLock
    {
        private final Lock lock = new ReentrantLock();
        private final Condition classGenerated = lock.newCondition();
        private volatile boolean released = false;

        public void await()
        {
            lock.lock();
            try
            {
                while (!released)
                    classGenerated.await();
            }
            catch (Exception ex)
            {
                throw undeclared(ex);
            }
            finally
            {
                lock.unlock();
            }
        }

        public void signal()
        {
            lock.lock();
            try
            {
                released = true;
                classGenerated.signalAll();
            }
            finally
            {
                lock.unlock();
            }
        }
    }

    private final Class<T> klas;
    private final MethodInterceptor interceptor;
    private volatile String name = "Proxy";


    /**
     * Creates a proxy builder that implements the non-final class <tt>klas</tt> with the callback
     * <tt>interceptor</tt>. This method is equivalent to the constructor but does not require a
     * redundant generic type argument.
     *
     * @param <T>
     * @param klas the class or interface to implement
     * @param interceptor the callback to implement with
     * @return
     */
    public static <T> ProxyBuilder<T> of(Class<T> klas, MethodInterceptor interceptor)
    {
        return new ProxyBuilder<T>(klas, interceptor);
    }

    /**
     * Creates a proxy builder that implements the non-final class <tt>klas</tt> with the callback
     * <tt>interceptor</tt>.
     *
     * @param klas the class or interface to implement
     * @param interceptor the callback to implement with
     */
    public ProxyBuilder(Class<T> klas, MethodInterceptor interceptor)
    {
        this.klas = klas;
        this.interceptor = interceptor;
    }

    /**
     * Sets the name to use in constructing the subclass name. The subclass will have the name
     * "package.OriginalClass$$Name$$", where "Name" is specified by this method. The default name
     * is "Proxy".
     *
     * @param name
     * @return this ProxyBuilder
     */
    public ProxyBuilder<T> setName(String name)
    {
        this.name = name;
        return this;
    }

    /**
     * Generates a new proxy instance without calling a constructor.
     *
     * @return new dynamic proxy instance
     */
    @SuppressWarnings("unchecked")
    public T create()
    {
        try
        {
            /* create an instance without invoking the constructor */
            T instance = (T)new ObjenesisStd().newInstance(getSubclass());

            /* set interceptor field */
            Field field = instance.getClass().getField(INTERCEPTOR_FIELD);
            field.set(instance, interceptor);

            return instance;
        }
        catch (Exception ex)
        {
            throw undeclared(ex);
        }
    }

    /**
     * Generates a new proxy instance calling the constructor whose signature matches the types of
     * <tt>args</tt>.
     *
     * @param args the constructor arguments
     * @return new dynamic proxy instance
     */
    public T create(Object... args)
    {
        Class<?>[] types = new Class[args.length];
        for (int i = 0 ; i < args.length ; i++)
        {
            Object arg = args[i];
            types[i] = arg == null ? Object.class : arg.getClass();
        }
        return create(types, args);
    }

    /**
     * Generates a new proxy instance calling the constructor whose signature matches <tt>types</tt>
     * with the arguments <tt>args</tt>.
     *
     * @param types
     * @param args
     * @return
     */
    public T create(Class<?>[] types, Object[] args)
    {
        if (types == null)
            /* recurse with inferred types */
            return create(args);

        Class<?>[] mungedTypes = new Class<?>[types.length + 1];
        Object[] mungedArgs = new Object[args.length + 1];

        mungedTypes[0] = MethodInterceptor.class;
        mungedArgs[0] = interceptor;

        for (int i = 0 ; i < args.length ; i++)
        {
            mungedTypes[i + 1] = types[i];
            mungedArgs[i + 1] = args[i];
        }

        try
        {
            return getSubclass().getConstructor(mungedTypes).newInstance(mungedArgs);
        }
        catch (InvocationTargetException ex)
        {
            throw undeclared(ex.getCause());
        }
        catch (Exception ex)
        {
            throw undeclared(ex);
        }
    }

    private String getSubclassInternal()
    {
        return getSubclassName().replace('.', '/');
    }

    private Type getSubclassType()
    {
        return Type.getType("L" + getSubclassInternal() + ";");
    }

    private String getSubclassName()
    {
        Package pkg = klas.getPackage();
        if (pkg != null && pkg.getName().startsWith("java."))
            /* we can't create classes in package java., so prepend something */
            return "_" + klas.getName() + "$$" + name + "$$";
        return klas.getName() + "$$" + name + "$$";
    }

    @SuppressWarnings("unchecked")
    private Class<? extends T> getSubclass()
    {
        for (Class<?> iface : klas.getInterfaces())
            if (iface == Proxy.class)
                /* base class is already a proxy */
                return klas;

        if ((klas.getModifiers() & 7) == 0 && klas.getPackage().getName().startsWith("java."))
            throw new IllegalStateException("Cannot proxy package protected classes in java.*");

        try
        {
            ClassLoader classLoader = klas.getClassLoader();
            if (classLoader == null)
                classLoader = ClassLoader.getSystemClassLoader();
            String subclassName = getSubclassName();


            /* synchronize by class name to ensure any class can only get generated once */
            GenerationLock lock = new GenerationLock();
            GenerationLock currentLock = generationLocks.putIfAbsent(subclassName, lock);
            if (currentLock != null)
            {
                lock = currentLock;
                lock.await();
            }

            try
            {
                return (Class<T>)classLoader.loadClass(subclassName);
            }
            catch (ClassNotFoundException ex)
            {
                try
                {
                    Class[] paramTypes = new Class[] { String.class, byte[].class, int.class, int.class };
                    Method method = ClassLoader.class.getDeclaredMethod("defineClass", paramTypes);
                    method.setAccessible(true);

                    byte[] bytes = getBytes();
                    return (Class<T>)method.invoke(classLoader, subclassName, bytes, 0, bytes.length);
                }
                finally
                {
                    lock.signal();
                }
            }
        }
        catch (Exception ex)
        {
            throw undeclared(ex);
        }
    }

    private byte[] getBytes()
    {
        final ClassWriter classWriter = new ClassWriter(ClassWriter.COMPUTE_MAXS);
        ClassVisitor visitor = classWriter;

        // visitor = new CheckClassAdapter( visitor );
        // visitor = new TraceClassVisitor( visitor, new PrintWriter( System.out ) );

        new ProxyWriter(this, visitor).write();

        return classWriter.toByteArray();
    }

    private static class ProxyWriter extends ClassAdapter
    {
        private final ProxyBuilder<?> proxy;

        public ProxyWriter(ProxyBuilder<?> proxy, ClassVisitor cv)
        {
            super(cv);
            this.proxy = proxy;
        }

        public void write()
        {
            String superclass = Type.getInternalName(proxy.klas);
            String proxyMarker = Type.getInternalName(Proxy.class);
            String[] interfaces;
            if (proxy.klas.isInterface())
            {
                interfaces = new String[] { superclass, proxyMarker };
                superclass = Type.getInternalName(Object.class);
            }
            else
            {
                interfaces = new String[] { proxyMarker };
            }

            visit(V1_5, ACC_PUBLIC, proxy.getSubclassInternal(), null, superclass, interfaces);

            /* add interceptor as a member variable */
            visitField(ACC_PUBLIC, INTERCEPTOR_FIELD, Type.getDescriptor(MethodInterceptor.class), null, null);

            /* add constructors that set interceptor and invoke super constructor */
            for (Constructor<?> ctor : proxy.klas.getDeclaredConstructors())
            {
                if ((ctor.getModifiers() & ACC_PRIVATE) == 0)
                    new ConstructorGenerator(ctor).write();
            }

            List<Method> methods = new ArrayList<Method>();
            addMethodsTo(methods, proxy.klas);
            if (proxy.klas.isInterface())
                addMethodsTo(methods, Object.class);

            /*
             * remove final methods (do this last in case method is defined multiple times with
             * different modifiers)
             */
            for (Iterator<Method> i = methods.iterator() ; i.hasNext() ;)
                if (Modifier.isFinal(i.next().getModifiers()))
                    i.remove();

            /* add methods that invoke interceptor */
            for (Method method : methods)
                new MethodGenerator(method).write();

            /* add methods that call super method */
            for (Method method : methods)
                if ((method.getModifiers() & (ACC_ABSTRACT | ACC_INTERFACE)) == 0)
                    new SuperGenerator(method).write();

            try
            {
                new GetInterceptorGenerator().write();
            }
            catch (NoSuchMethodException ex)
            {
                throw undeclared(ex);
            }

            visitEnd();
        }

        private static void addMethodsTo(List<Method> methods, Class<?> klas)
        {
            /* add all non-private, non-static methods that haven't been added already */
            for (Method method : klas.getDeclaredMethods())
            {
                if ((method.getModifiers() & (ACC_PRIVATE | ACC_STATIC)) == 0)
                {
                    boolean alreadyAdded = false;
                    for (Method each : methods)
                    {
                        if (each.getName().equals(method.getName())
                                && Arrays.equals(each.getParameterTypes(), method.getParameterTypes()))
                        {
                            alreadyAdded = true;
                            break;
                        }
                    }
                    if (!alreadyAdded)
                        methods.add(method);
                }
            }

            Class<?> superclass = klas.getSuperclass();
            if (superclass != null)
                addMethodsTo(methods, superclass);

            for (Class<?> iface : klas.getInterfaces())
                addMethodsTo(methods, iface);
        }

        private static String getDescriptor(Constructor<?> ctor, boolean addInterceptor)
        {
            StringBuilder s = new StringBuilder();
            s.append("(");

            if (addInterceptor)
                /* add method interceptor as first parameter */
                s.append(Type.getDescriptor(MethodInterceptor.class));

            for (Class<?> c : ctor.getParameterTypes())
                s.append(Type.getDescriptor(c));

            s.append(")V");
            return s.toString();
        }

        private class ConstructorGenerator extends GeneratorAdapter
        {
            private final Constructor<?> ctor;

            public ConstructorGenerator(Constructor<?> ctor)
            {
                this(ctor, getDescriptor(ctor, true));
            }

            private ConstructorGenerator(Constructor<?> ctor, String desc)
            {
                super(cv.visitMethod(ACC_PUBLIC, "<init>", desc, null, null), ACC_PUBLIC, "<init>", desc);
                this.ctor = ctor;
            }

            public void write()
            {
                mv.visitCode();

                /* set interceptor */
                loadThis();
                loadArg(0);
                putField(proxy.getSubclassType(), INTERCEPTOR_FIELD, Type.getType(MethodInterceptor.class));

                /* invoke super class constructor */
                loadThis();
                loadArgs(1, ctor.getParameterTypes().length);
                invokeConstructor(Type.getType(ctor.getDeclaringClass()),
                                  new org.objectweb.asm.commons.Method("<init>", getDescriptor(ctor, false)));

                returnValue();
                mv.visitMaxs(0, 0);
                mv.visitEnd();
            }
        }

        /**
         * Generates a special method for each proxied method that delegates to the original
         * ("super") implementation.
         */
        private class SuperGenerator extends GeneratorAdapter
        {
            private final Method method;

            public SuperGenerator(Method method)
            {
                this(method, method.getName() + SUPER_METHOD_SUFFIX, Type.getMethodDescriptor(method));
            }

            private SuperGenerator(Method method, String name, String desc)
            {
                super(cv.visitMethod(ACC_PUBLIC, name, desc, null, null), ACC_PUBLIC, name, desc);
                this.method = method;
            }

            public void write()
            {
                mv.visitCode();

                /* call super class implementation */
                loadThis();
                loadArgs();

                String owner = Type.getInternalName(method.getDeclaringClass());
                String desc = Type.getMethodDescriptor(method);
                mv.visitMethodInsn(INVOKESPECIAL, owner, method.getName(), desc);

                returnValue();

                mv.visitMaxs(0, 0);
                mv.visitEnd();
            }
        }

        /**
         * Generates a special <tt>$getInterceptor</tt> method which implements the <tt>Proxy</tt>
         * interface.
         */
        private class GetInterceptorGenerator extends NormalMethodGenerator
        {
            public GetInterceptorGenerator() throws NoSuchMethodException
            {
                this(Proxy.class.getMethod("$getInterceptor"));
            }

            private GetInterceptorGenerator(Method method)
            {
                this(method, method.getName(), Type.getMethodDescriptor(method));
            }

            private GetInterceptorGenerator(Method method, String name, String desc)
            {
                super(cv.visitMethod(ACC_PUBLIC, name, desc, null, null), ACC_PUBLIC, name, desc);
            }

            public void write()
            {
                try
                {
                    mv.visitCode();

                    /* return interceptor */
                    Type interceptorType = Type.getType(MethodInterceptor.class);
                    loadThis();
                    getField(proxy.getSubclassType(), INTERCEPTOR_FIELD, interceptorType);
                    checkCast(interceptorType);
                    returnValue();

                    mv.visitMaxs(0, 0);
                    mv.visitEnd();
                }
                catch (Exception ex)
                {
                    throw undeclared(ex);
                }
            }
        }

        /**
         * Generates proxy methods that invoke the MethodInterceptor for this class.
         */
        private class MethodGenerator extends NormalMethodGenerator
        {
            private final Method method;

            public MethodGenerator(Method method)
            {
                this(method, method.getName(), Type.getMethodDescriptor(method));
            }

            private MethodGenerator(Method method, String name, String desc)
            {
                super(cv.visitMethod(ACC_PUBLIC, name, desc, null, null), ACC_PUBLIC, name, desc);
                this.method = method;
            }

            public void write()
            {
                try
                {
                    mv.visitCode();

                    /* put interceptor on the stack */
                    Type interceptorType = Type.getType(MethodInterceptor.class);
                    loadThis();
                    getField(proxy.getSubclassType(), INTERCEPTOR_FIELD, interceptorType);

                    /* get the Method object for this method */

                    /* push new InvocationContext */
                    Type invocationContextType = Type.getType(InvocationContext.class);
                    newInstance(invocationContextType);
                    dup();

                    push(Type.getType(method.getDeclaringClass()));
                    push(method.getName());

                    /* create a class array with the method parameter types */
                    Class<?>[] parameterTypes = method.getParameterTypes();
                    push(parameterTypes.length);
                    newArray(Type.getType(Class.class));

                    for (int i = 0 ; i < parameterTypes.length ; i++)
                    {
                        dup();
                        push(i);
                        pushBoxType(Type.getType(parameterTypes[i]));
                        arrayStore(Type.getType(Class.class));
                    }

                    invokeVirtual("java/lang/Class", "getDeclaredMethod",
                                  "(Ljava/lang/String;[Ljava/lang/Class;)Ljava/lang/reflect/Method;");

                    /* push object and args constructor arguments */
                    loadThis();
                    loadArgArray();

                    /* call InvocationContext constructor */
                    invokeSpecial("org/mockfu/InvocationContext", "<init>",
                                  "(Ljava/lang/reflect/Method;Ljava/lang/Object;[Ljava/lang/Object;)V");

                    /* invoke interceptor */
                    invokeInterface("org/mockfu/MethodInterceptor", "intercept",
                                    "(Lorg/mockfu/InvocationContext;)Ljava/lang/Object;");

                    Type returnType = Type.getReturnType(method);
                    Label nonNullLabel = new Label();
                    dup();

                    ifNonNull(nonNullLabel);
                    /* if return value is null don't unbox, just return default value */
                    pop();
                    pushDefault(returnType);
                    returnValue();

                    mark(nonNullLabel);
                    unbox(returnType);
                    returnValue();

                    mv.visitMaxs(0, 0);
                    mv.visitEnd();
                }
                catch (Exception ex)
                {
                    throw undeclared(ex);
                }
            }


        }
    }

}
